from tree_sitter_languages import get_language, get_parser
from collections import namedtuple

Tag = namedtuple("Tag", "rel_fname fname line name kind".split())
language = get_language('python')
parser = get_parser('python')

def get_ast_tree_from_code(code):
    tree = parser.parse(code.encode())
    root = tree.root_node
    print("展示整个语法树的结构")
    print(root.sexp())
    # return tree 还是return root
    return tree

def query_ast_tree(tree, query_scm):
    tree_root_node = tree.root_node
    doc_str_query = language.query(query_scm)
    captures = doc_str_query.captures(tree_root_node)
    return captures

def parse_tree_captures_to_tags(rel_fname, fname, captures):
    for node, tag in captures:

        print(f"node:{node}")
        print(f"tag:{tag}")
        name = node.text.decode("utf-8")
        if tag.startswith("name.definition."):
            kind = "def"
        elif tag.startswith("name.reference."):
            kind = "ref"
        else:
            continue
        result = Tag(
            rel_fname=rel_fname,
            fname=fname,
            # 源代码
            name=name,
            # 参考scm文件，一般方法的调用是reference
            kind=kind,
            line=node.start_point[0],
        )
        yield result